import numpy as np
#初始化矩阵W d=4
W = np.array([[4, 3, 2, 1],
              [2, 2, 2, 2],
              [1, 3, 4, 2],
              [0, 1, 2, 3]])

#矩阵的维度
d = W.shape[0]
#秩
r = 2
#随机初始化A和B
np.random.seed(666)
#A服从标准正态分布
A = np.random.randn(d,r)
B = np.zeros((r,d))
print(A)
print(B)

#定义超参数
lr = 0.01
epochs = 1000

#定义损失函数
def loss_function(W, A, B):
    W_approx = A @ B
    #均方差损失函数
    loss = np.linalg.norm(W - W_approx,'fro')**2
    return loss

#定义梯度下降法
def descent(W, A, B,epochs):
    loss_history = []
    for i in range(epochs):
        #计算梯度
        W_approx = A @ B
        #计算损失函数关于矩阵A和B的梯度
        gd_A = -2 * (W - W_approx) @ B.T
        gd_B = -2 * A.T @ (W - W_approx)
        #更新A和B
        A -= lr *gd_A
        B -= lr *gd_B
        #计算当前的损失
        loss = loss_function(W, A, B)
        loss_history.append(loss)
        if i % 100 == 0:
            print(f"Epoch {i}: Loss {loss}")
    return A, B, loss_history

A, B, loss_history = descent(W, A, B, epochs)
print(A)
print(B)

import matplotlib.pyplot as plt

plt.figure(figsize=(8,6))
plt.plot(loss_history,label="loss")
plt.xlabel("epochs")
plt.ylabel("loss")
plt.grid(True)
plt.show()
